import torch
import torchvision
import onnx
import sys, os
sys.path.insert(0, sys.path[0]+"/../")
import copy
from modules import loadckpts, BuildSegmentor
from myutil.mixconfig import SEGMENTOR_CFG
if __name__ == "__main__":
    segmentor = BuildSegmentor(segmentor_cfg=copy.deepcopy(SEGMENTOR_CFG), mode='TEST')
    ckpt = loadckpts("fcn_resnet50os16_voc/epoch_3.pth")
    segmentor.load_state_dict(ckpt['model'])
    segmentor.to("cuda:0")
    # torch.save(ckpt["model"], "a.pt")
    dummy_input = torch.randn(1,3,512,512,device="cuda")
    # model = torchvision.models.alexnet(pretrained=True).cuda()
    torch.onnx.export(segmentor, dummy_input, "fcn.onnx", verbose=True)
